#!/usr/bin/env python3
"""
MCP服务器主入口文件
支持HTTP、Stdio和WebSocket三种连接方式
"""

# 抑制librosa的pkg_resources弃用警告
import warnings
warnings.filterwarnings('ignore', category=UserWarning, module='librosa')

import asyncio
import sys
import argparse
from pathlib import Path
from typing import Dict, Any

from loguru import logger

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from src.config import load_config
from src.mcp.server import MCPServer
from src.core.context_orchestrator import ContextOrchestrator
from src.database import DatabaseManager


async def start_http_server(config: Dict[str, Any]):
    """启动HTTP服务器"""
    try:
        from fastapi import FastAPI
        from fastapi.middleware.cors import CORSMiddleware
        import uvicorn
        import json
        
        # 初始化上下文编排器
        orchestrator = ContextOrchestrator(config)
        await orchestrator.initialize()
        
        # 创建MCP服务器
        mcp_server = MCPServer(orchestrator, config)
        await mcp_server.start()
        
        # 创建FastAPI应用
        app = FastAPI(
            title="MCP Server",
            description="Model Context Protocol Server with HTTP and SSE support",
            version="1.0.0"
        )
        
        # 添加CORS中间件
        app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )
        
        @app.get("/")
        async def root():
            """根路径"""
            return {"message": "MCP Server is running", "version": "1.0.0"}
        
        @app.get("/health")
        async def health_check():
            """健康检查"""
            return await mcp_server.health_check()
        
        @app.get("/stats")
        async def server_stats():
            """服务器统计信息"""
            return await mcp_server.get_server_stats()
        
        @app.post("/mcp")
        async def handle_mcp_request(request_data: dict):
            """处理MCP请求"""
            try:
                message = json.dumps(request_data, ensure_ascii=False)
                response = await mcp_server.handle_message(message)
                return json.loads(response)
            except Exception as e:
                logger.error(f"处理MCP请求失败: {e}")
                return {
                    "jsonrpc": "2.0",
                    "error": {
                        "code": -32603,
                        "message": "Internal error",
                        "data": str(e)
                    },
                    "id": request_data.get("id")
                }
        
        @app.get("/mcp")
        async def mcp_sse_endpoint():
            """MCP Server-Sent Events端点"""
            from fastapi.responses import StreamingResponse
            
            async def event_stream():
                """SSE事件流"""
                connection_id = f"sse_{id(event_stream)}"
                
                try:
                    # 连接到MCP服务器
                    await mcp_server.connect_client(connection_id, {
                        'transport': 'sse',
                        'endpoint': '/mcp'
                    })
                    
                    # 发送初始连接事件
                    yield f"data: {{\"type\": \"connected\", \"connection_id\": \"{connection_id}\"}}\n\n"
                    
                    # 保持连接活跃
                    while True:
                        # 发送心跳
                        yield f"data: {{\"type\": \"heartbeat\", \"timestamp\": {asyncio.get_event_loop().time()}}}\n\n"
                        await asyncio.sleep(30)  # 每30秒发送一次心跳
                        
                except asyncio.CancelledError:
                    logger.info(f"SSE连接已取消: {connection_id}")
                except Exception as e:
                    logger.error(f"SSE连接错误: {e}")
                    yield f"data: {{\"type\": \"error\", \"message\": \"{str(e)}\"}}\n\n"
                finally:
                    # 断开连接
                    try:
                        await mcp_server.disconnect_client(connection_id)
                    except:
                        pass
            
            return StreamingResponse(
                event_stream(),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                    "Access-Control-Allow-Origin": "*",
                    "Access-Control-Allow-Headers": "*",
                }
            )
        
        # 获取HTTP配置
        http_config = config.get('mcp', {}).get('http', {})
        host = http_config.get('host', 'localhost')
        port = http_config.get('port', 8000)
        
        # 启动服务器
        server_config = uvicorn.Config(
            app,
            host=host,
            port=port,
            log_level="info",
            access_log=True
        )
        
        server = uvicorn.Server(server_config)
        logger.info(f"启动HTTP服务器: http://{host}:{port}")
        logger.info(f"MCP端点: http://{host}:{port}/mcp")
        logger.info(f"SSE端点: http://{host}:{port}/mcp (GET)")
        logger.info(f"健康检查: http://{host}:{port}/health")
        
        await server.serve()
        
    except ImportError as e:
        logger.error(f"缺少依赖: {e}")
        logger.error("请安装: pip install fastapi uvicorn")
        sys.exit(1)
    except Exception as e:
        logger.error(f"启动HTTP服务器失败: {e}")
        raise


async def start_stdio_server(config: Dict[str, Any]):
    """启动Stdio服务器"""
    from .stdio_handler import run_stdio_server
    
    # 初始化上下文编排器
    orchestrator = ContextOrchestrator(config)
    await orchestrator.initialize()
    
    # 创建MCP服务器
    mcp_server = MCPServer(orchestrator, config)
    
    # 在stdio模式下避免日志输出
    # logger.info("启动Stdio MCP服务器")
    await run_stdio_server(mcp_server, config)


async def start_websocket_server(config: Dict[str, Any]):
    """启动WebSocket服务器"""
    from .websocket_handler import WebSocketHandler
    
    # 初始化上下文编排器
    orchestrator = ContextOrchestrator(config)
    await orchestrator.initialize()
    
    # 创建MCP服务器
    mcp_server = MCPServer(orchestrator, config)
    
    # 创建WebSocket处理器
    ws_handler = WebSocketHandler(mcp_server, config)
    
    ws_config = config.get('mcp', {}).get('websocket', {})
    host = ws_config.get('host', 'localhost')
    port = ws_config.get('port', 8765)
    
    logger.info(f"启动WebSocket服务器: ws://{host}:{port}")
    await ws_handler.start()
    
    # 保持服务器运行
    try:
        while ws_handler.is_running:
            await asyncio.sleep(1)
    except KeyboardInterrupt:
        logger.info("收到中断信号，正在停止WebSocket服务器...")
        await ws_handler.stop()


def setup_logging(config: Dict[str, Any], mode: str = 'http'):
    """配置日志"""
    log_config = config.get('logging', {})
    
    # 移除默认处理器
    logger.remove()
    
    # 在stdio模式下，为了确保MCP协议兼容性，只使用文件日志
    if mode == 'stdio':
        # 只添加文件处理器，避免任何控制台输出
        if log_config.get('file_path'):
            log_file = Path(log_config['file_path'])
            log_file.parent.mkdir(parents=True, exist_ok=True)
            
            logger.add(
                log_file,
                level=log_config.get('level', 'INFO'),
                format=log_config.get('format', '{time} | {level} | {message}'),
                rotation=log_config.get('rotation', '1 day'),
                retention=log_config.get('retention', '30 days'),
                compression=log_config.get('compression', 'gz')
            )
        else:
            # 如果没有配置文件日志，则完全禁用日志
            logger.add(lambda _: None, level="CRITICAL")
    else:
        # 非stdio模式，正常配置日志
        # 添加控制台处理器
        logger.add(
            sys.stderr,
            level=log_config.get('level', 'INFO'),
            format=log_config.get('format', '{time} | {level} | {message}'),
            colorize=True
        )
        
        # 添加文件处理器
        if log_config.get('file_path'):
            log_file = Path(log_config['file_path'])
            log_file.parent.mkdir(parents=True, exist_ok=True)
            
            logger.add(
                log_file,
                level=log_config.get('level', 'INFO'),
                format=log_config.get('format', '{time} | {level} | {message}'),
                rotation=log_config.get('rotation', '1 day'),
                retention=log_config.get('retention', '30 days'),
                compression=log_config.get('compression', 'gz')
            )


async def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='MCP服务器')
    parser.add_argument(
        '--http', 
        action='store_true',
        help='启动HTTP服务器模式'
    )
    parser.add_argument(
        '--stdio', 
        action='store_true',
        help='启动Stdio服务器模式'
    )
    parser.add_argument(
        '--websocket', 
        action='store_true',
        help='启动WebSocket服务器模式'
    )
    parser.add_argument(
        '--port', 
        type=int,
        help='服务器端口号'
    )
    parser.add_argument(
        '--host', 
        type=str,
        default='localhost',
        help='服务器主机地址'
    )
    parser.add_argument(
        '--config', 
        default='config',
        help='配置文件目录路径'
    )
    
    args = parser.parse_args()
    
    # 确定启动模式
    if args.http:
        mode = 'http'
    elif args.stdio:
        mode = 'stdio'
    elif args.websocket:
        mode = 'websocket'
    else:
        # 默认HTTP模式
        mode = 'http'
    
    # 在stdio模式下，在加载配置之前就静默所有日志，避免任何stdout/stderr污染
    if mode == 'stdio':
        try:
            logger.remove()
        except Exception:
            pass
        # 使用空接收器吞掉所有日志
        logger.add(lambda _: None, level="CRITICAL")
    
    # 加载配置
    try:
        config = load_config(args.config)
        
        # 覆盖配置中的端口和主机
        if args.port:
            if mode == 'http':
                config.setdefault('mcp', {}).setdefault('http', {})['port'] = args.port
            elif mode == 'websocket':
                config.setdefault('mcp', {}).setdefault('websocket', {})['port'] = args.port
        
        if args.host:
            if mode == 'http':
                config.setdefault('mcp', {}).setdefault('http', {})['host'] = args.host
            elif mode == 'websocket':
                config.setdefault('mcp', {}).setdefault('websocket', {})['host'] = args.host
        
        setup_logging(config, mode)
        if mode != 'stdio':  # 在stdio模式下避免任何日志输出
            logger.info(f"配置加载成功，启动模式: {mode}")
    except Exception as e:
        # 在stdio模式下，不能使用print输出到stdout，必须输出到stderr
        if 'mode' in locals() and mode == 'stdio':
            import traceback
            tb = traceback.format_exc()
            sys.stderr.write(f"配置加载失败: {e}\n{tb}\n")
        else:
            logger.error(f"配置加载失败: {e}")
        sys.exit(1)
    
    # 根据模式启动相应服务器
    try:
        if mode == 'http':
            await start_http_server(config)
        elif mode == 'stdio':
            await start_stdio_server(config)
        elif mode == 'websocket':
            await start_websocket_server(config)
    except KeyboardInterrupt:
        if mode != 'stdio':
            logger.info("收到中断信号，正在关闭服务器...")
        # stdio 模式下避免任何输出
    except Exception as e:
        if mode == 'stdio':
            import traceback
            tb = traceback.format_exc()
            # 仅向 stderr 输出，避免污染 stdout
            sys.stderr.write(f"服务器启动失败: {e}\n{tb}\n")
        else:
            logger.error(f"服务器启动失败: {e}")
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())